# -*- coding: utf-8 -*-
import uuid
from typing import List

from sqlalchemy import func, and_, or_, not_
from sqlalchemy.orm import selectinload

from db.session import with_session
from db.models.knowledge import KnowledgeBase, File2Knowledge, RecommendQuestion
from db.models.files import QAFile, DualControlFile, EmergencyFile, AccidentCaseFile
from db.models.attributes import FileAttributeAssociation, ApplicableObject, KnowledgeType, ApplicableArea, ElementClassification


""" =========================== KnowledgeBase ================================ """


# 查找 根据 知识库id 查找知识库
@with_session
def db_get_kb_by_id(session, kb_id):
    return session.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id, KnowledgeBase.is_delete.is_(False)).first()


@with_session
def db_get_kb_by_name(session, kb_name):
    return session.query(KnowledgeBase).filter(KnowledgeBase.name == kb_name, KnowledgeBase.is_delete.is_(False)).first()


# 查找 获取所有 知识库 列表
@with_session
def db_get_all_kbs(session, kb_type: int = None):
    if kb_type is None:
        return session.query(KnowledgeBase).filter(KnowledgeBase.is_delete.is_(False)).all()
    else:
        return session.query(KnowledgeBase).filter(KnowledgeBase.kb_type == str(kb_type), KnowledgeBase.is_delete.is_(False)).all()


# 获取 用户的 所有知识库
@with_session
def db_get_kbs_by_user(session, user, user_group=[], kb_type: int = 1):
    params = KnowledgeBase.owner == user.id
    for i in user_group:
        if params is None:
            params = func.json_contains(KnowledgeBase.group_ids, str(i))
        else:
            params = params | func.json_contains(KnowledgeBase.group_ids, str(i))
    # kb_type 字段存储为字符串，确保比较类型一致
    kbs = list(
        session.query(KnowledgeBase)
        .filter(
            KnowledgeBase.is_delete.is_(False),
            KnowledgeBase.kb_type == str(kb_type)
        )
        .filter(params)
        .all()
    )
    return kbs


@with_session
def get_user_knowledge_bases(session, user_id: int, is_superuser: bool, user_group):
    if is_superuser:
        return session.query(KnowledgeBase).filter(KnowledgeBase.is_delete.is_(False)).all()
    else:
        params = KnowledgeBase.owner == user_id
        # 注意：这里需要根据实际情况调整用户组查询逻辑
        # 假设 get_user_group_ids 是一个单独的函数
        for group_id in user_group:
            if params is None:
                params = func.json_contains(KnowledgeBase.group_ids, str(group_id))
            else:
                params = params | func.json_contains(KnowledgeBase.group_ids, str(group_id))
        return session.query(KnowledgeBase).filter(KnowledgeBase.is_delete.is_(False)).filter(params).all()


# 增加知识库
@with_session
def db_create_kb(session, name, user_id, group_ids, kb_type):
    kb = session.query(KnowledgeBase).filter(KnowledgeBase.name == name).first()
    if kb:
        kb.is_delete = False
        kb.name = name
        kb.group_ids = group_ids
        kb.kb_type = str(kb_type)
        session.commit()
        return kb
    kb = KnowledgeBase(name=name, owner=user_id, group_ids=group_ids, kb_type=str(kb_type))
    session.add(kb)
    session.commit()
    return kb


@with_session
def db_update_kb_name(session, kb_id, new_name):
    kb = session.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id).first()
    if not kb:
        return None
    kb.name = new_name
    session.commit()
    return kb


# 删除知识库
@with_session
def db_delete_kb_by_id(session, kb_id):
    kb = session.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id, KnowledgeBase.is_delete.is_(False)).first()
    if not kb:
        return None
    kb.is_delete = True
    session.commit()
    return kb


@with_session
def db_update_kb_groups(session, kb_id, new_group_ids):
    kb = session.query(KnowledgeBase).filter(KnowledgeBase.id == kb_id, KnowledgeBase.is_delete.is_(False)).first()
    if not kb:
        return None
    kb.group_ids = new_group_ids
    session.commit()
    return kb


# 查找 在对应知识库中根据 文件名称和四个属性查找
@with_session
def db_get_files_in_kb_by_params(
        session,
        knowledge_id: uuid.UUID,
        file_name: str = "",
        applicable_object: List[int] = None,
        knowledge_type: List[int] = None,
        applicable_area: List[int] = None,
        element_classification: List[int] = None,
        offset: int = 0,
        limit: int = 50
) -> List[QAFile]:
    kb = db_get_kb_by_id(knowledge_id)
    model = {
        '1': QAFile,
        '2': DualControlFile,
        '3': EmergencyFile,
        '4': AccidentCaseFile,
    }.get(str(kb.kb_type))
    if not model:
        return [], 0

    query = session.query(model).join(
        File2Knowledge,
        or_(
            File2Knowledge.file_id == func.BIN_TO_UUID(model.id),
            File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(model.id), '-', ''),
            File2Knowledge.file_id == func.HEX(model.id),
            File2Knowledge.file_id == func.LOWER(func.HEX(model.id))
        )
    ).filter(
        File2Knowledge.knowledge_id == knowledge_id,
        model.is_delete.is_(False)
    )

    if model is QAFile:
        query = query.options(selectinload(QAFile.attributes))

    if file_name:
        query = query.filter(model.name.ilike(f"%{file_name}%"))

    if model is QAFile:
        attr_filters = {
            "applicable_object": applicable_object or [],
            "knowledge_type": knowledge_type or [],
            "applicable_area": applicable_area or [],
            "element_classification": element_classification or []
        }

        subqueries = [
            QAFile.id.in_(
                session.query(FileAttributeAssociation.file_id)
                .filter(
                    FileAttributeAssociation.attribute_type == attr_type,
                    FileAttributeAssociation.attribute_id.in_(attr_ids)
                )
                .subquery()
            )
            for attr_type, attr_ids in attr_filters.items()
            if attr_ids
        ]

        if subqueries:
            query = query.filter(and_(*subqueries))

    total = query.order_by(None).count()
    items = query.order_by(model.create_time.desc()).offset(offset).limit(limit).all()
    return items, total


@with_session
def db_get_files_not_in_kb_by_params(
        session,
        knowledge_id: uuid.UUID,
        file_name: str = "",
        applicable_object: List[int] = None,
        knowledge_type: List[int] = None,
        applicable_area: List[int] = None,
        element_classification: List[int] = None,
        offset: int = 0,
        limit: int = 50
) -> List[QAFile]:
    kb = db_get_kb_by_id(knowledge_id)
    model = {
        '1': QAFile,
        '2': DualControlFile,
        '3': EmergencyFile,
        '4': AccidentCaseFile,
    }.get(str(kb.kb_type))
    if not model:
        return [], 0

    query = session.query(model).filter(model.is_delete.is_(False))

    if model is QAFile:
        query = query.options(selectinload(QAFile.attributes))

    if file_name:
        query = query.filter(model.name.ilike(f"%{file_name}%"))

    if model is QAFile:
        attr_filters = {
            "applicable_object": applicable_object or [],
            "knowledge_type": knowledge_type or [],
            "applicable_area": applicable_area or [],
            "element_classification": element_classification or []
        }

        subqueries = [
            QAFile.id.in_(
                session.query(FileAttributeAssociation.file_id)
                .filter(
                    FileAttributeAssociation.attribute_type == attr_type,
                    FileAttributeAssociation.attribute_id.in_(attr_ids)
                )
                .subquery()
            )
            for attr_type, attr_ids in attr_filters.items()
            if attr_ids
        ]

        if subqueries:
            query = query.filter(and_(*subqueries))

    file_id_condition = or_(
        File2Knowledge.file_id == func.BIN_TO_UUID(model.id),
        File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(model.id), '-', ''),
        File2Knowledge.file_id == func.HEX(model.id),
        File2Knowledge.file_id == func.LOWER(func.HEX(model.id))
    )
    exclude_condition = session.query(File2Knowledge.id).filter(
        File2Knowledge.knowledge_id == knowledge_id,
        file_id_condition
    ).exists()
    query = query.filter(~exclude_condition)

    total = query.count()
    items = query.order_by(model.create_time.desc()).offset(offset).limit(limit).all()
    return items, total


@with_session
def db_get_all_files_by_params(
        session,
        file_name: str = "",
        applicable_object: List[int] = None,
        knowledge_type: List[int] = None,
        applicable_area: List[int] = None,
        element_classification: List[int] = None,
        offset: int = 0,
        limit: int = 50
) -> List[QAFile]:
    if applicable_object is None:
        applicable_object = []
    if knowledge_type is None:
        knowledge_type = []
    if applicable_area is None:
        applicable_area = []
    if element_classification is None:
        element_classification = []

    query = session.query(QAFile) \
        .options(selectinload(QAFile.attributes)) \
        .filter(QAFile.is_delete.is_(False))

    if file_name:
        query = query.filter(QAFile.name.ilike(f"%{file_name}%"))

    subqueries = []
    attr_filters = [
        ("applicable_object", applicable_object),
        ("knowledge_type", knowledge_type),
        ("applicable_area", applicable_area),
        ("element_classification", element_classification)
    ]

    for attr_type, attr_ids in attr_filters:
        if attr_ids:
            subquery = session.query(FileAttributeAssociation.file_id).filter(
                FileAttributeAssociation.attribute_type == attr_type,
                FileAttributeAssociation.attribute_id.in_(attr_ids)
            ).subquery()
            subqueries.append(QAFile.id.in_(subquery))

    if subqueries:
        query = query.filter(and_(*subqueries))

    total = query.count()
    items = query.order_by(QAFile.create_time.desc()).offset(offset).limit(limit).all()
    return items, total

""" ================================= RecommendQuestion ================================ """

@with_session
def db_get_recommend_question(session, knowledge_id):
    return session.query(RecommendQuestion).filter(RecommendQuestion.knowledge_id == knowledge_id, RecommendQuestion.is_delete.is_(False)).all()


@with_session
def db_add_recommend_question(session, knowledge_id, question):
    recommend_question = RecommendQuestion(knowledge_id=knowledge_id, question=question)
    session.add(recommend_question)
    session.commit()
    return recommend_question

@with_session
def db_update_recommend_question(session, recommend_question_id, question):
    recommend_question = session.query(RecommendQuestion).get(recommend_question_id)
    recommend_question.question = question
    session.commit()
    return recommend_question

@with_session
def db_delete_recommend_question(session, recommend_question_id):
    recommend_question = session.query(RecommendQuestion).get(recommend_question_id)
    recommend_question.is_delete = True
    session.commit()
    return recommend_question


""" ================================= File2Kb ================================ """


# 查找 根据 知识库id、文件id 查找关联
@with_session
def db_get_asso_by_kb_file(session, kb_id, file_id):
    # 确保file_id是字符串格式
    file_id_str = str(file_id) if file_id is not None else ""
    return session.query(File2Knowledge) \
        .filter(File2Knowledge.knowledge_id == kb_id, File2Knowledge.file_id == file_id_str) \
        .all()


@with_session
def db_get_files_in_kb(session, kb_id, kb_type: int = 1, *, name_like: str = None, offset: int = 0, limit: int = 50, all: bool = False):
    """获取知识库中所有未删除的文件，kb_type=1返回 QAFile，kb_type=2 返回 DualControlFile"""
    if str(kb_type) == '1':
        query = session.query(QAFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(QAFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(QAFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(QAFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(QAFile.id))
                )
            ) \
            .options(selectinload(QAFile.attributes)) \
            .filter(File2Knowledge.knowledge_id == kb_id, QAFile.is_delete.is_(False))
        if name_like:
            query = query.filter(QAFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        if all:
            return query.all(), total
        items = query.order_by(QAFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '2':
        query = session.query(DualControlFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(DualControlFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(DualControlFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(DualControlFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(DualControlFile.id))
                )
            ) \
            .filter(File2Knowledge.knowledge_id == kb_id, DualControlFile.is_delete.is_(False))
        if name_like:
            query = query.filter(DualControlFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(DualControlFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '3':
        # 应急预案文件
        query = session.query(EmergencyFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(EmergencyFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(EmergencyFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(EmergencyFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(EmergencyFile.id))
                )
            ) \
            .filter(File2Knowledge.knowledge_id == kb_id, EmergencyFile.is_delete.is_(False))
        if name_like:
            query = query.filter(EmergencyFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(EmergencyFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '4':
        # 事故案例文件
        query = session.query(AccidentCaseFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(AccidentCaseFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(AccidentCaseFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(AccidentCaseFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(AccidentCaseFile.id))
                )
            ) \
            .filter(File2Knowledge.knowledge_id == kb_id, AccidentCaseFile.is_delete.is_(False))
        if name_like:
            query = query.filter(AccidentCaseFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(AccidentCaseFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    return [], 0


@with_session
def db_get_files_not_in_kb(session, kb_id, kb_type: int = 1, *, name_like: str = None, offset: int = 0, limit: int = 50):
    """获取知识库中所有未删除的文件，kb_type=1返回 QAFile，kb_type=2 返回 DualControlFile"""
    if str(kb_type) == '1':
        query = session.query(QAFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(QAFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(QAFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(QAFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(QAFile.id))
                )
            ) \
            .options(selectinload(QAFile.attributes)) \
            .filter(not_(File2Knowledge.knowledge_id == kb_id), QAFile.is_delete.is_(False))
        if name_like:
            query = query.filter(QAFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(QAFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '2':
        query = session.query(DualControlFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(DualControlFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(DualControlFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(DualControlFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(DualControlFile.id))
                )
            ) \
            .filter(not_(File2Knowledge.knowledge_id == kb_id), DualControlFile.is_delete.is_(False))
        if name_like:
            query = query.filter(DualControlFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(DualControlFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '3':
        # 应急预案文件
        query = session.query(EmergencyFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(EmergencyFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(EmergencyFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(EmergencyFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(EmergencyFile.id))
                )
            ) \
            .filter(not_(File2Knowledge.knowledge_id == kb_id), EmergencyFile.is_delete.is_(False))
        if name_like:
            query = query.filter(EmergencyFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(EmergencyFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    if str(kb_type) == '4':
        # 事故案例文件
        query = session.query(AccidentCaseFile) \
            .join(
                File2Knowledge,
                or_(
                    File2Knowledge.file_id == func.BIN_TO_UUID(AccidentCaseFile.id),
                    File2Knowledge.file_id == func.REPLACE(func.BIN_TO_UUID(AccidentCaseFile.id), '-', ''),
                    File2Knowledge.file_id == func.HEX(AccidentCaseFile.id),
                    File2Knowledge.file_id == func.LOWER(func.HEX(AccidentCaseFile.id))
                )
            ) \
            .filter(not_(File2Knowledge.knowledge_id == kb_id), AccidentCaseFile.is_delete.is_(False))
        if name_like:
            query = query.filter(AccidentCaseFile.name.ilike(f"%{name_like}%"))
        total = query.count()
        items = query.order_by(AccidentCaseFile.create_time.desc()).offset(offset).limit(limit).all()
        return items, total
    return [], 0


# 删除 关联
@with_session
def db_delete_kb_file_asso(session, asso: List):
    for i in asso:
        session.delete(i)
    session.commit()
    return True


@with_session
def db_delete_files_in_kb(session, kb_id):
    try:
        asso_list = session.query(File2Knowledge) \
            .filter(File2Knowledge.knowledge_id == kb_id) \
            .delete(synchronize_session=False)  # 禁用同步以提升性能

        # for i in asso_list:
        #     session.delete(i)

        session.commit()
        return True
    except Exception as e:
        session.rollback()
        # 需要导入logger或使用其他日志记录方式
        # print(f"删除知识库 {kb_id} 的文件关联失败: {e}")
        return False


@with_session
def db_add_kb_file_asso(session, kb_id, file_id):
    if kb_id is None:
        return None
    if not isinstance(kb_id, uuid.UUID):
        kb_id = uuid.UUID(kb_id)
    # 确保file_id是字符串格式
    file_id_str = str(file_id) if file_id is not None else ""
    asso = File2Knowledge(knowledge_id=kb_id, file_id=file_id_str)
    session.add(asso)
    session.commit()
    return asso
